package com.tuniu.agents.common.agent;

import com.tuniu.agents.advisor.AgentChatMemoryAdvisor;
import com.tuniu.agents.agent.*;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.message.MessageLogger;
import com.tuniu.agents.util.ChatResponseUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

public class SelectionAgent extends OptionsAgent {

    protected final MessageLogger messageLogger;

    protected final Map<String, String> toolAnalysisMap;

    protected final AgentManager agentManager;

    /**
     * Analyst：
     * The user previously raised these requests:
     * 1: XXX
     * 2: XXX
     * The request of the current user is:
     * XXX
     */
    public static final String MESSAGE_TEMPLATE = """
            {{user_message_his}}
            The original request of the current user is:
            {{user_message_current}}
            """;

    public SelectionAgent(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, MessageLogger messageLogger, Map<String, String> toolAnalysisMap, String... tools) {
        super(agentManagerGroup, new AgentOptions(agentId, promptTemplateId, List.of(tools)));
        this.messageLogger = messageLogger;
        this.toolAnalysisMap = toolAnalysisMap;
        this.agentManager = agentManagerGroup.getAgentManager();
    }

    @Override
    public Flux<String> stream(String conversationId, List<Message> message, Map<String,Object> context) {
        ChatResponse chatResponse = this.callResponse(conversationId, message, context);
        String toolName = ChatResponseUtils.getFirstToolName(chatResponse);
        String toolResponse = ChatResponseUtils.getFirstTextMessage(chatResponse);
        if (!StringUtils.hasText(toolName)) {
            return Flux.just(ChatResponseUtils.getFirstTextMessage(chatResponse));
        } else {
            Agent analysisAgent = agentManager.getAgent(toolAnalysisMap.get(toolName));
            Assert.notNull(analysisAgent, "Failed to get local Agent By toolName:" + toolName);
            return analysisAgent.stream(conversationId, toolResponse);
        }
    }

    @Override
    public Flux<ChatResponse> streamResponse(String conversationId, List<Message> message, Map<String,Object> context) {
        ChatResponse chatResponse = this.callResponse(conversationId, message, context);
        String toolName = ChatResponseUtils.getFirstToolName(chatResponse);
        String toolResponse = ChatResponseUtils.getFirstTextMessage(chatResponse);
        if (!StringUtils.hasText(toolName)) {
            return Flux.just(chatResponse);
        } else {
            Agent analysisAgent = agentManager.getAgent(toolAnalysisMap.get(toolName));
            Assert.notNull(analysisAgent, "Failed to get local Agent By toolName:" + toolName);
            return analysisAgent.streamResponse(conversationId, List.of(new UserMessage(toolResponse)), context);
        }
    }

    @Override
    public ChatClient.ChatClientRequestSpec chatClientRequestSpec(String conversationId, List<Message> messages, Map<String,Object> context) {
        return chatClient.prompt()
                .system(s -> s.params(chatEnvironmentMemory.get(conversationId)))
                .messages(messages)
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId).
                        param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, super.agentOptions.getAdvisorParam().get(CHAT_MEMORY_RETRIEVE_SIZE_KEY)).
                        param(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_TEMPLATE, SelectionAgent.MESSAGE_TEMPLATE)
                        .param(AgentsConstants.AgentsAdvisorKey.ASSISTANT_MESSAGE_HISTORY_NEED, "false")
                        .param(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_HISTORY_NEED, "false")
                        .param(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY, this.agentId));
    }


}
